/*
* ////////////////////////////////////////////////
*
* 开发者：AZPEP
* 友情链接：
*	-	B站：		零分的考卷			https://space.bilibili.com/257391258
*	-	Github：	AZPEP-BiliBili		https://github.com/AZPEP-BiliBili
*	-	Gitee：		AZPEP_BiliBili		https://gitee.com/AZPEP_BiliBili
*
* CopyRight AZPEP
*
* ////////////////////////////////////////////////
*/

#include "Global_Include.h"
#include "Entry.h"
#include <format>

/// <summary>
/// 初始化程序
/// </summary>
void Init();
std::wstring Dword2Wstring(DWORD d);
/// <summary>
///  获取磁盘上所有空闲的盘符
/// </summary>
/// <returns>成功GetLastError()=0，失败GetLastError()为非零值</returns>
std::vector<WCHAR> AstInitGetUsedLetter();

/*
/// <summary>
/// 程序真正入口
/// </summary>
void RealProgram() {
	Init();
	RealRuntime();


}
*/
/// <summary>
/// 初始化程序
/// </summary>
void Init() {
	_wsetlocale(0, L"chs");

	/* 开启日志服务 */

	LogHandle = AstLog_StartService(
		L"\"AZPEP SysRi Tool\" - 主程序"
	);

	//检查日志服务状态
	switch (AstLog_CheckHandleAvaiable(LogHandle)) {
		//成功
	case 0:
		break;

		//失败，报错退出
	case 1:
	case 2:
		AstThrowError_ThrowError_NonWriteLog(
			ERROR_INIT_START_LOG_SERVICE_FAIL,
			*new std::wstring{ L"[INIT]启动日志服务失败！WinErr代码：" + Dword2Wstring(GetLastError()) }
		);
		AstThrowError_EmergencyExit_NonWriteLog(1);
	}
	AstLog_WriteMessage(
		LogHandle,
		std::format(L"日志文件：{}", AstLog_GetLogFilePath(LogHandle))
	);
	AstLog_WriteWarning(LogHandle, L"[INIT]！！！程序初始化开始！！！");

	/* 初始化ThrowError*/
	AstLog_WriteMessage(LogHandle, L"[INIT][THROWERROR]开始启用日志异常记录服务");
	if (!AstThrowError_Init(LogHandle)) {
		//初始化失败，报错退出
		AstThrowError_ThrowError_NonWriteLog(
			ERROR_INIT_START_LOG_SERVICE_FAIL,
			*new std::wstring{ L"[INIT][THROWERROR]启动日志异常记录服务失败！WinErr代码：" + Dword2Wstring(GetLastError()) }
		);
		AstThrowError_EmergencyExit_NonWriteLog(1);
	}
	AstLog_WriteMessage(LogHandle, L"[INIT][THROWERROR]日志异常记录服务已启用！");


	/* 获取固件类型 */
	AstLog_WriteMessage(LogHandle, L"[INIT][FIRMWARE]开始检查计算机固件类型");
	try {
		FIRMWARE_TYPE firmwareType;
		if (GetFirmwareType(&firmwareType))
		{
			if (firmwareType != FIRMWARE_TYPE::FirmwareTypeUefi)
			{
				AstLog_WriteMessage(LogHandle, L"[INIT][FIRMWARE]电脑固件类型：非UEFI");
				AstThrowError_ThrowError(
					ERROR_INIT_NO_UEFI_FIRMWARE,
					*new std::wstring{ L"[INIT][FIRMWARE]此Windows操作系统不是以UEFI方式启动，请使用UEFI方式重启电脑后再运行本软件" }
				);
				AstThrowError_EmergencyExit(1);
			}
		}
	}
	catch (std::wstring ErrInfo) {
		//捕获到异常，报错
		AstThrowError_ThrowError(
			ERROR_INIT_GET_FIRMWARE_TYPE_FAIL,
			*new std::wstring{ L"[INIT][FIRMWARE]获取启动固件类型失败！\n错误信息：" + ErrInfo }
		);
		AstThrowError_EmergencyExit(1);
	}
	AstLog_WriteMessage(LogHandle, L"[INIT][FIRMWARE]电脑固件类型：UEFI");


	/* 获取EFI分区表 */
	AstLog_WriteMessage(LogHandle, L"[INIT][ESP]开始获取EFI分区");
	AstMountEsp_UnMountEfiPartitionW();
	DWORD dStatus = AstMountEsp_MountEfiPartitionW(&EspLetterList);

	//获取失败，报错退出
	if (dStatus != 0) {
		AstThrowError_ThrowError(
			ERROR_INIT_GET_EFI_PARTITION_FAIL,
			*new std::wstring{ L"[INIT][ESP]获取EFI分区失败！WinErr代码：" + Dword2Wstring(dStatus) }
		);
		AstThrowError_EmergencyExit(1);
	}
	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][ESP]已找到EFI分区{}个", EspLetterList.size()));
	for (auto iter = EspLetterList.begin(); iter != EspLetterList.end(); iter++) {
		WCHAR w = (*iter);
		std::wstringstream Path;
		Path << w << L":\\";
		std::wstring PathW = Path.str();
		LPCWSTR PathW2 = PathW.c_str();
		//获取磁盘类型
		UINT driveType = GetDriveTypeW(PathW2);
		//不是固定分区，弹出警告，报错
		if (driveType != DRIVE_FIXED && driveType != DRIVE_NO_ROOT_DIR) {
			AstMountEsp_UnmountDriveW((*iter));
			EspLetterList.erase(iter);
			AstLog_WriteMessage(LogHandle, std::format(L"[INIT][ESP]驱动器“{}”的类型为{},不是“Fixed”，不可用于重装系统！", (*iter), driveType));

			iter--;
		}
		else if (driveType == DRIVE_NO_ROOT_DIR) {
			AstThrowError_ThrowError(
				ERROR_INIT_GET_EFI_PARTITION_FAIL,
				*new std::wstring{ L"[INIT][ESP]获取EFI驱动器类型失败！WinErr代码：" + Dword2Wstring(dStatus) }
			);
			AstMountEsp_UnMountEfiPartitionW();
			AstThrowError_EmergencyExit(1);
		}

	}
	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][ESP]可供重装使用的分区共{}个", EspLetterList.size()));

	//没有ESP分区，报错退出
	if (EspLetterList.size() == 0) {
		AstThrowError_ThrowError(
			ERROR_INIT_GET_EFI_PARTITION_FAIL,
			*new std::wstring{ L"[INIT][ESP]计算机中没有可供重装使用的EFI分区！" }
		);
		//卸载所有EFI分区
		AstMountEsp_UnMountEfiPartitionW();
		AstThrowError_EmergencyExit(1);
	}

	//收集EFI驱动器信息
	AstLog_WriteMessage(LogHandle, L"[INIT][ESP]开始收集EFI分区信息");
	std::vector<SimpleEfiPartitionInformationW*> EspInfomationListTemp;
	dStatus = AstMontEsp_GetAllEfiPartitionInformationW(&EspInfomationListTemp);
	//获取失败，报错退出
	if (dStatus != 0) {
		AstThrowError_ThrowError(
			ERROR_INIT_GET_EFI_PARTITION_FAIL,
			*new std::wstring{ L"[INIT][EFI]获取EFI分区信息表失败！WinErr代码：" + Dword2Wstring(dStatus) }
		);
		//卸载所有EFI分区
		AstMountEsp_UnMountEfiPartitionW();

		AstThrowError_EmergencyExit(1);
	}

	//开始对应信息
	for (auto iter : EspInfomationListTemp) {
		EspInfomationList.push_back(new EfiPartitionMap());
		bool correspond = false;
		int Count = EspInfomationList.size() - 1;

		//对应信息
		for (auto letter : EspLetterList) {
			if (EspInfomationListTemp[Count]->Letter == letter) {
				EspInfomationList[Count]->Letter = letter;
				EspInfomationList[Count]->PartitionInfo = EspInfomationListTemp[Count];

				correspond = true;

				break;
			}
		}

		if (!correspond) {
			AstThrowError_ThrowError(
				ERROR_INIT_GET_EFI_PARTITION_FAIL,
				std::format(L"[INIT][ESP]无法对应EFI驱动器“{}”的信息", EspInfomationListTemp[Count]->Letter)
			);

			//卸载EFI分区
			AstMountEsp_UnMountEfiPartitionW();
			AstThrowError_EmergencyExit(1);
		}

	}
	AstLog_WriteMessage(LogHandle, L"[INIT][ESP]信息对应完成！");

#ifdef _DEBUG
	///输出信息
	AstLog_WriteDebug(LogHandle, L"[INIT][ESP]\n[INIT][ESP]以下是EFI驱动器信息：\n[INIT][ESP]--------------------");

	for (auto iter : EspInfomationList) {
		AstLog_WriteDebug(LogHandle, std::format(L"[INIT][ESP]盘符：{}", iter->Letter));
		AstLog_WriteDebug(LogHandle, std::format(L"[INIT][ESP]硬盘序号：{}", iter->PartitionInfo->DiskNumber));
		AstLog_WriteDebug(LogHandle, std::format(L"[INIT][ESP]分区序号：{}", iter->PartitionInfo->PartitionNumber));
		AstLog_WriteDebug(LogHandle, std::format(L"[INIT][ESP]分区路径：{}", iter->PartitionInfo->PartitionPath));
		AstLog_WriteDebug(LogHandle, std::format(L"[INIT][ESP]分区卷路径：{}", iter->PartitionInfo->PartitionVolumePath));
		AstLog_WriteDebug(LogHandle, L"[INIT][ESP]--------------------\n");
	}
#endif

	/* 获取系统盘信息 */
	AstLog_WriteMessage(LogHandle, L"[INIT][PARTITION]开始获取系统盘信息");

	//获取环境变量
	LPWSTR Var = (LPWSTR)malloc(MAX_PATH);
	ZeroMemory(Var, MAX_PATH);
	dStatus = GetEnvironmentVariable(
		L"SystemDrive",
		Var,
		MAX_PATH
	);
	//获取失败，报错退出
	if (dStatus == 0) {
		AstThrowError_ThrowError(
			ERROR_INIT_GET_EFI_PARTITION_FAIL,
			*new std::wstring{ L"[INIT][PARTITION]获取操作系统分区盘符失败！WinErr代码：" + Dword2Wstring(dStatus) }
		);
		//卸载所有EFI分区
		AstMountEsp_UnMountEfiPartitionW();

		AstThrowError_EmergencyExit(1);
	}
	std::wstring SystemRoot{ Var };
	SystemPartitionLetter = SystemRoot[0];

	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][PARTITION]操作系统分区盘符为：{}", SystemPartitionLetter));

	//检查系统驱动器是否为GPT分区
	BOOL bStatus;
	dStatus = AstMountEsp_IsPartitionGptW(
		&SystemPartitionLetter,
		&bStatus
	);
	//获取失败，退出
	if (dStatus != 0) {
		AstThrowError_ThrowError(
			ERROR_INIT_GET_SYSTEM_PARTITION_STYLE_FAIL,
			std::format(L"[INIT][PARTITION]获取系统分区“{}”的分区样式失败！WinErr代码：{}", SystemPartitionLetter, dStatus)
		);
		AstMountEsp_UnMountEfiPartitionW();
		AstThrowError_EmergencyExit(1);
	}

	//不是GPT分区，报错退出
	if (!bStatus) {
		AstLog_WriteMessage(LogHandle, std::format(L"系统分区“{}”分区样式：MBR", SystemPartitionLetter));

		AstThrowError_ThrowError(
			ERROR_INIT_SYSTEM_PARTITION_ISNT_A_GPT_PARTITION,
			std::format(L"[INIT][PARTITION]系统分区“{}”不是一个GPT分区，请联系电脑脑维修人员将分区类型转换为GPT分区后重试", SystemPartitionLetter)
		);
		AstMountEsp_UnMountEfiPartitionW();
		AstThrowError_EmergencyExit(1);
	}
	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][PARTITION]系统分区“{}”分区样式：GPT", SystemPartitionLetter));

	// 拼凑磁盘路径
	std::wstringstream wss;
	wss << SystemPartitionLetter << L":";


	/* 获取系统总大小 */
	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][PARTITION]开始获取系统分区“{}”总空间大小", SystemPartitionLetter));
	_ULARGE_INTEGER a;
	ULARGE_INTEGER nFreeBytesAvailable;
	ULARGE_INTEGER nTotalNumberOfBytes;
	ULARGE_INTEGER nTotalNumberOfFreeBytes;

	//利用GetDiskFreeSpaceEx函数获得指定磁盘信息
	if (!GetDiskFreeSpaceEx(wss.str().c_str(),
		&nFreeBytesAvailable,
		&nTotalNumberOfBytes,
		&nTotalNumberOfFreeBytes)
		) {
		//获取失败，报错
		AstThrowError_ThrowError(
			ERROR_INIT_SYSTEM_PARTITION_ISNT_A_GPT_PARTITION,
			std::format(L"[INIT][PARTITION]获取系统分区“{}”失败，WinErr代码：{}", SystemPartitionLetter, GetLastError()));
		AstMountEsp_UnMountEfiPartitionW();
		AstThrowError_EmergencyExit(1);
	}
	//计算分区大小
	SystemPartitionTotalSpace = *new unsigned long long(nTotalNumberOfBytes.QuadPart);
	AstLog_WriteMessage(LogHandle, std::format(L"[INIT][PARTITION]操作系统分区“{}”总空间：{} 字节", SystemPartitionLetter, SystemPartitionTotalSpace));

	/* 获取其他分区总大小 */
	AstLog_WriteMessage(LogHandle, L"[INIT][PARTITION]开始获取其他固定分区信息");
	// 获取已用盘符
	auto LetterList = AstInitGetUsedLetter();

	for (auto Letter : LetterList) {
		//是ESP分区，跳过
		if (std::count(EspLetterList.begin(), EspLetterList.end(), Letter) != 0) {
			AstLog_WriteWarning(LogHandle, std::format(L"[INIT][PARTITION]驱动器“{}”是EFI分区，跳过", Letter));
			continue;
		}

		//不是GPT分区，跳过
		dStatus = AstMountEsp_IsPartitionGptW(
			&Letter,
			&bStatus
		);
		//获取失败，退出
		if (dStatus != 0) {
			AstThrowError_ThrowError(
				ERROR_INIT_GET_SYSTEM_PARTITION_STYLE_FAIL,
				std::format(L"[INIT][PARTITION]获取分区“{}”的分区样式失败！WinErr代码：{}", SystemPartitionLetter, dStatus)
			);
			AstMountEsp_UnMountEfiPartitionW();
			AstThrowError_EmergencyExit(1);
		}

		if (!bStatus) {
			AstLog_WriteWarning(LogHandle, std::format(L"[INIT][PARTITION]驱动器“{}”不是GPT分区，跳过", Letter));
			continue;
		}

		std::wstringstream Path;
		Path << Letter << L":\\";
		std::wstring PathW = Path.str();
		LPCWSTR PathW2 = PathW.c_str();
		//获取磁盘类型
		UINT driveType = GetDriveTypeW(PathW2);
		if (driveType != DRIVE_FIXED) {
			AstLog_WriteWarning(LogHandle, std::format(L"[INIT][PARTITION]驱动器“{}”不是固定分区，跳过", Letter));
			continue;
		}
		//开始获取
		//利用GetDiskFreeSpaceEx函数获得指定磁盘信息
		if (!GetDiskFreeSpaceEx(PathW.c_str(),
			&nFreeBytesAvailable,
			&nTotalNumberOfBytes,
			&nTotalNumberOfFreeBytes)
			) {
			//获取失败，报错
			AstThrowError_ThrowError(
				ERROR_INIT_SYSTEM_PARTITION_ISNT_A_GPT_PARTITION,
				std::format(L"[INIT][PARTITION]获取分区“{}”信息失败，WinErr代码：{}", SystemPartitionLetter, GetLastError()));
			AstMountEsp_UnMountEfiPartitionW();
			AstThrowError_EmergencyExit(1);
		}

		//计算总大小
		unsigned long long PartitionTotalSpace = *new unsigned long long(nTotalNumberOfBytes.QuadPart);
		PartitionInfoW PIW;
		PIW.Letter = *new WCHAR(Letter);
		PIW.TotalSpace = *new unsigned long long(PartitionTotalSpace);
		PIWL.push_back(*new PartitionInfoW(PIW));

		AstLog_WriteMessage(LogHandle, std::format(L"[INIT][PARTITION]驱动器“{}”已加载，总大小：{} 字节", PIW.Letter, PIW.TotalSpace));
	}

	AstLog_WriteWarning(LogHandle, L"[INIT]！！！程序初始化结束！！！");

	return;
}

std::wstring Dword2Wstring(DWORD d) {
	std::wstringstream ss;
	ss << d;
	return ss.str();
}

/// <summary>
///  获取磁盘上所有已使用的盘符
/// </summary>
/// <returns>成功GetLastError()=0，失败GetLastError()为非零值</returns>
std::vector<WCHAR> AstInitGetUsedLetter() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'C'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << L"\\\\.\\" << wc << L":";

		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);

		if (hDevice != INVALID_HANDLE_VALUE) {
			CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}

	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}
